package grouping;

import org.apache.lucene.analysis.core.WhitespaceAnalyzer;
import org.apache.lucene.document.*;
import org.apache.lucene.index.*;
import org.apache.lucene.queries.function.ValueSource;
import org.apache.lucene.queries.function.valuesource.BytesRefFieldSource;
import org.apache.lucene.search.IndexSearcher;
import org.apache.lucene.search.Query;
import org.apache.lucene.search.Sort;
import org.apache.lucene.search.TermQuery;
import org.apache.lucene.search.grouping.*;
import org.apache.lucene.search.similarities.BM25Similarity;
import org.apache.lucene.store.MMapDirectory;
import org.apache.lucene.util.BytesRef;
import org.apache.lucene.util.mutable.MutableValue;
import org.apache.lucene.util.mutable.MutableValueStr;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;

import java.io.IOException;
import java.nio.file.Paths;
import java.util.ArrayList;
import java.util.Collection;
import java.util.HashMap;
import java.util.List;

import static org.apache.lucene.tests.util.LuceneTestCase.random;
import static org.junit.jupiter.api.Assertions.*;

public class GroupingDemo {
//    @Test
//    public void singlePassForSearchMethodTwo() {
//        // Per search:
//        GroupingSearch groupingSearch = new GroupingSearch(groupEndDocs);
//        groupingSearch.setGroupSort(groupSort);
//        groupingSearch.setIncludeScores(needsScores);
//        TermQuery query = new TermQuery(new Term("content", searchTerm));
//        TopGroups groupsResult = groupingSearch.search(indexSearcher, query, groupOffset, groupLimit);
//
//        // Render groupsResult...
//    }
//
//    @Test
//    public void singlePassForSearchMethodOne() {
//        Query groupEndDocs = new TermQuery(new Term("groupEnd", "x"));
//        BlockGroupingCollector c = new BlockGroupingCollector(groupSort, groupOffset + topNGroups, needsScores, groupEndDocs);
//        s.search(new TermQuery(new Term("content", searchTerm)), c);
//        TopGroups groupsResult = c.getTopGroups(withinGroupSort, groupOffset, docOffset, docOffset + docsPerGroup, fillFields);
//
//        // Render groupsResult...
//    }
//
//    @Test
//    public void singlePassForIndexing() {
//        // Create Documents from your source:
//        List<Document> oneGroup = ...;
//
//        Field groupEndField = new Field("groupEnd", "x", Field.Store.NO, Field.Index.NOT_ANALYZED);
//        groupEndField.setIndexOptions(IndexOptions.DOCS_ONLY);
//        groupEndField.setOmitNorms(true);
//        oneGroup.get(oneGroup.size() - 1).add(groupEndField);
//
//        // You can also use writer.updateDocuments(); just be sure you
//        // replace an entire previous doc block with this new one.  For
//        // example, each group could have a "groupID" field, with the same
//        // value for all docs in this group:
//        writer.addDocuments(oneGroup);
//    }

    MMapDirectory mMapDirectory = null;
    IndexWriter w =null;

    @BeforeEach
    public void setUp(){


        String uri = "F:\\integration\\search\\lucene-data\\grouping";
        try {
            mMapDirectory = new MMapDirectory(Paths.get(uri));
            IndexWriterConfig indexWriterConfig = new IndexWriterConfig();
            indexWriterConfig.setOpenMode(IndexWriterConfig.OpenMode.CREATE);

            WhitespaceAnalyzer whitespaceAnalyzer = new WhitespaceAnalyzer(30);
            w = new IndexWriter(mMapDirectory, new IndexWriterConfig(whitespaceAnalyzer));

        } catch (IOException e) {
            throw new RuntimeException("mmap failed");
        }

    }
    @AfterEach
    public void destroyDown(){
        try {
            mMapDirectory.close();
        } catch (IOException e) {
            e.printStackTrace();
        }

    }

    private void addGroupField(Document doc, String groupField, String value) {
        doc.add(new SortedDocValuesField(groupField, new BytesRef(value)));
    }
    @SuppressWarnings({"unchecked", "rawtypes"})
    private <T> TopGroupsCollector<T> createSecondPassCollector(
            FirstPassGroupingCollector firstPassGroupingCollector,
            Sort groupSort,
            Sort sortWithinGroup,
            int groupOffset,
            int maxDocsPerGroup,
            boolean getMaxScores)
            throws IOException {

        Collection<SearchGroup<T>> searchGroups = firstPassGroupingCollector.getTopGroups(groupOffset);
        return new TopGroupsCollector<>(
                firstPassGroupingCollector.getGroupSelector(),
                searchGroups,
                groupSort,
                sortWithinGroup,
                maxDocsPerGroup,
                getMaxScores);
    }

    // Basically converts searchGroups from MutableValue to BytesRef if grouping by ValueSource
    @SuppressWarnings("unchecked")
    private TopGroupsCollector<?> createSecondPassCollector(
            FirstPassGroupingCollector<?> firstPassGroupingCollector,
            String groupField,
            Collection<SearchGroup<BytesRef>> searchGroups,
            Sort groupSort,
            Sort sortWithinGroup,
            int maxDocsPerGroup,
            boolean getMaxScores)
            throws IOException {
        if (firstPassGroupingCollector
                .getGroupSelector()
                .getClass()
                .isAssignableFrom(TermGroupSelector.class)) {
            GroupSelector<BytesRef> selector =
                    (GroupSelector<BytesRef>) firstPassGroupingCollector.getGroupSelector();
            return new TopGroupsCollector<>(
                    selector, searchGroups, groupSort, sortWithinGroup, maxDocsPerGroup, getMaxScores);
        } else {
            ValueSource vs = new BytesRefFieldSource(groupField);
            List<SearchGroup<MutableValue>> mvalSearchGroups = new ArrayList<>(searchGroups.size());
            for (SearchGroup<BytesRef> mergedTopGroup : searchGroups) {
                SearchGroup<MutableValue> sg = new SearchGroup<>();
                MutableValueStr groupValue = new MutableValueStr();
                if (mergedTopGroup.groupValue != null) {
                    groupValue.value.copyBytes(mergedTopGroup.groupValue);
                } else {
                    groupValue.exists = false;
                }
                sg.groupValue = groupValue;
                sg.sortValues = mergedTopGroup.sortValues;
                mvalSearchGroups.add(sg);
            }
            ValueSourceGroupSelector selector = new ValueSourceGroupSelector(vs, new HashMap<>());
            return new TopGroupsCollector<>(
                    selector, mvalSearchGroups, groupSort, sortWithinGroup, maxDocsPerGroup, getMaxScores);
        }
    }
    private FirstPassGroupingCollector<?> createRandomFirstPassCollector(
            String groupField, Sort groupSort, int topDocs) throws IOException {
        if (random().nextBoolean()) {
            ValueSource vs = new BytesRefFieldSource(groupField);
            return new FirstPassGroupingCollector<>(
                    new ValueSourceGroupSelector(vs, new HashMap<>()), groupSort, topDocs);
        } else {
            return new FirstPassGroupingCollector<>(
                    new TermGroupSelector(groupField), groupSort, topDocs);
        }
    }
    private void compareGroupValue(String expected, GroupDocs<?> group) {
        if (expected == null) {
            if (group.groupValue == null) {
                return;
            } else if (group.groupValue.getClass().isAssignableFrom(MutableValueStr.class)) {
                return;
            } else if (((BytesRef) group.groupValue).length == 0) {
                return;
            }
            fail();
        }

        if (group.groupValue.getClass().isAssignableFrom(BytesRef.class)) {
            assertEquals(new BytesRef(expected), group.groupValue);
        } else if (group.groupValue.getClass().isAssignableFrom(MutableValueStr.class)) {
            MutableValueStr v = new MutableValueStr();
            v.value.copyChars(expected);
            assertEquals(v, group.groupValue);
        } else {
            fail();
        }
    }
    @Test
    public void index() throws IOException {
        FieldType customType = new FieldType();
        customType.setStored(true);
        String groupField = "author";

        // 0
        Document doc = new Document();
        addGroupField(doc, groupField, "author1");
        doc.add(new TextField("content", "random text", Field.Store.YES));
        doc.add(new Field("id", "1", customType));
        w.addDocument(doc);

        // 1
        doc = new Document();
        addGroupField(doc, groupField, "author1");
        doc.add(new TextField("content", "some more random text", Field.Store.YES));
        doc.add(new Field("id", "2", customType));
        w.addDocument(doc);

        // 2
        doc = new Document();
        addGroupField(doc, groupField, "author1");
        doc.add(new TextField("content", "some more random textual data", Field.Store.YES));
        doc.add(new Field("id", "3", customType));
        w.addDocument(doc);

        // 3
        doc = new Document();
        addGroupField(doc, groupField, "author2");
        doc.add(new TextField("content", "some random text", Field.Store.YES));
        doc.add(new Field("id", "4", customType));
        w.addDocument(doc);

        // 4
        doc = new Document();
        addGroupField(doc, groupField, "author3");
        doc.add(new TextField("content", "some more random text", Field.Store.YES));
        doc.add(new Field("id", "5", customType));
        w.addDocument(doc);

        // 5
        doc = new Document();
        addGroupField(doc, groupField, "author3");
        doc.add(new TextField("content", "random", Field.Store.YES));
        doc.add(new Field("id", "6", customType));
        w.addDocument(doc);

        // 6 -- no author field
        doc = new Document();
        doc.add(new TextField("content", "random word stuck in alot of other text", Field.Store.YES));
        doc.add(new Field("id", "6", customType));
        w.addDocument(doc);
        w.close();;
    }
    @Test
    public void twoPass() throws IOException {
        String groupField="author";
        DirectoryReader reader = DirectoryReader.open(mMapDirectory);

        IndexSearcher indexSearcher = new IndexSearcher(reader);
        // This test relies on the fact that longer fields produce lower scores
        indexSearcher.setSimilarity(new BM25Similarity());


        final Sort groupSort = Sort.RELEVANCE;

        final FirstPassGroupingCollector<?> c1 =
                createRandomFirstPassCollector(groupField, groupSort, 10);
        indexSearcher.search(new TermQuery(new Term("content", "random")), c1);
// from the group field 'author' to group ,
        final TopGroupsCollector<?> c2 =
                createSecondPassCollector(c1, groupSort, Sort.RELEVANCE, 0, 5, true);
        indexSearcher.search(new TermQuery(new Term("content", "random")), c2);

        final TopGroups<?> groups = c2.getTopGroups(0);
        assertFalse(Float.isNaN(groups.maxScore));

        assertEquals(7, groups.totalHitCount);
        assertEquals(7, groups.totalGroupedHitCount);
        assertEquals(4, groups.groups.length);

        // relevance order: 5, 0, 3, 4, 1, 2, 6

        // the later a document is added the higher this docId
        // value
        GroupDocs<?> group = groups.groups[0];
        compareGroupValue("author3", group);
        assertEquals(2, group.scoreDocs.length);
        assertEquals(5, group.scoreDocs[0].doc);
        assertEquals(4, group.scoreDocs[1].doc);
        assertTrue(group.scoreDocs[0].score > group.scoreDocs[1].score);

        group = groups.groups[1];
        compareGroupValue("author1", group);
        assertEquals(3, group.scoreDocs.length);
        assertEquals(0, group.scoreDocs[0].doc);
        assertEquals(1, group.scoreDocs[1].doc);
        assertEquals(2, group.scoreDocs[2].doc);
        assertTrue(group.scoreDocs[0].score >= group.scoreDocs[1].score);
        assertTrue(group.scoreDocs[1].score >= group.scoreDocs[2].score);

        group = groups.groups[2];
        compareGroupValue("author2", group);
        assertEquals(1, group.scoreDocs.length);
        assertEquals(3, group.scoreDocs[0].doc);

        group = groups.groups[3];
        compareGroupValue(null, group);
        assertEquals(1, group.scoreDocs.length);
        assertEquals(6, group.scoreDocs[0].doc);

        indexSearcher.getIndexReader().close();
        w.close();
    }

}

